import numpy as np
import pandas as pd
import os
from core.constant import *
import copy

import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib
import seaborn as sns
from conf import conf

import squarify

# 设置中文字体为SimHei
mpl.rcParams['font.sans-serif'] = ['SimHei']
# 设置负号正常显示
mpl.rcParams['axes.unicode_minus'] = False

matplotlib.use("Agg")


class PlanPLT:
    def __init__(self, parent):
        """
        致开发者：
            如果您看到了这里，那么我想您已经在尝试自主开发因子分析方案。作为一个数学专业出身的理科生，
        我承认该框架的代码结构并不优秀，这个缺点会让您在编写分析方案或者做其他扩展时遇到一些不必要的麻
        烦。
            截止上线，该项目的代码编写和商业化架构均由我独立完成。我向往一个漂亮的开源的金融数据分析
        框架，也有信心实现这个目标。但在当前处境下，尽快变现是一个很现实的问题。所以，一些开发任务会被
        搁置。
            在获得一定的现金流之后，我将会组建专业的技术团队，开发下一代DDQuant。届时，DDQuant将不
        依赖第三方软件，能够在自己的生态中加入金融数据库、全品种实盘交易接口、AI智能开发等高阶功能。金
        融数据分析部分仍将保持开源，同时具备足够优秀的代码结构和开发支持。
            相对于市面上现有的量化交易框架，DDQuant的数据分析能力仍然是非常有竞争力的！
            感谢您的支持！

        :param parent:
        """
        self.parent = parent

    def create_and_save_plot(self, dataframe, plot_type, factor_ls, result_dir_path, symbol):
        """
        生成并保存图表。
        :param dataframe: 数据框架
        :param plot_type: 图表类型 ('Profile', 'Density', 'Frequency')
        :param factor_ls: 因子列表
        :param result_dir_path: 结果目录路径
        :param symbol: 标志符号
        """
        plt.figure(figsize=(15, 6))
        base_name = f"{symbol}-{plot_type}.jpg"
        fig_path = os.path.join(result_dir_path, base_name)

        if plot_type == 'Profile':
            plt.title("自然分布图")
            plt.xlabel("时间")
            plt.ylabel("因子值")
            dataframe = dataframe.dropna()
            for i, column in enumerate(factor_ls):
                # val_sum = sum(abs(dataframe[column]))
                # 将数据缩放至同一区间
                scale_y = ((dataframe[column] - dataframe[column].min()) /
                           (dataframe[column].max() - dataframe[column].min()))
                plt.plot(range(len(scale_y.index)), scale_y.values, color=plt.cm.tab20(i / len(factor_ls)),
                         label=f'{column}')
            legend = plt.legend()
            # for handle in legend.legendHandles:
            #     handle.set_sizes([25])

        elif plot_type == 'Density':
            plt.title("顺序分布图")
            plt.xlabel("自增序号")
            plt.ylabel("归一因子值")
            dataframe = dataframe.dropna()
            sorted_df = pd.DataFrame(index=range(len(dataframe)))
            for col in factor_ls:
                factor_sr = dataframe[col]
                factor_sr = (factor_sr - factor_sr.min()) / (factor_sr.max() - factor_sr.min())
                sorted_df[col] = factor_sr.sort_values().values
            for i, column in enumerate(factor_ls):
                draw_y = sorted_df[column]
                plt.plot(range(len(draw_y.index)), draw_y.values, color=plt.cm.tab20(i / len(factor_ls)),
                         label=f'{column}')
            legend = plt.legend()
            # for handle in legend.legendHandles:
            #     handle.set_sizes([25])

        elif plot_type == 'Frequency':
            plt.title("频率分布直方图")
            plt.xlabel("归一因子值")
            plt.ylabel("群组大小")
            dataframe = dataframe.dropna()
            num_bins = 40
            for i, column in enumerate(factor_ls):
                scale_y = ((dataframe[column] - dataframe[column].min()) /
                           (dataframe[column].max() - dataframe[column].min()))
                # plt.hist(scale_y, bins=30, alpha=0.5, step=True, edgecolor='black', linewidth=1.5, color=plt.cm.tab20(i / len(factor_ls)),
                #          label=f'{column}')
                counts, bin_edges = np.histogram(scale_y, bins=num_bins, range=(0, 1))
                plt.plot(bin_edges[1:], counts, color=plt.cm.tab20(i / len(factor_ls)),
                         label=f'{column}')
            legend = plt.legend()
            # for handle in legend.legendHandles:
            #     handle.set_sizes([25])

        self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
        plt.close()

        # 返回图像的 Markdown 语法
        return f"![{symbol} {plot_type}](./{base_name})"

    def gen_save_fig(self, stocks, result_dir_path):
        base_name = f"选股云图.jpg"
        fig_path = os.path.join(result_dir_path, base_name)
        # 获取代码名称映射df
        support_data_path = os.path.join(self.parent.master.file_manager.support_path,
                                         conf.Framework.code_name_info_csv.value)
        support_data_df = self.parent.master.file_manager.read_csv(support_data_path)
        # 计算方块大小
        sizes = [abs(value) for value in stocks.values()]

        sum_sizes = np.sum(sizes)
        tap_value = sum_sizes * 0.01

        draw_sizes = []
        draw_stocks_dc = {}
        label_ls = []
        table_data = [["代码", "名称", "值", "显示"]]
        for i, (symbol, value) in enumerate(stocks.items()):
            if sizes[i] >= tap_value:
                draw_sizes.append(abs(value))
                draw_stocks_dc[symbol] = value
                name = support_data_df[support_data_df["code"] == symbol]["name"].values[0]
                label_ls.append(
                    '\n'.join([symbol + ' ' + name,
                               str(round(value, 3))]))
                table_data.append([symbol, name, round(list(stocks.values())[i], 4), "显示"])
            else:
                name = support_data_df[support_data_df["code"] == symbol]["name"].values[0]
                table_data.append([symbol, name, round(list(stocks.values())[i], 4), "不显示"])
        # 对table_data按照从大到小排序
        table_data[1:] = sorted(table_data[1:], key=lambda row: row[2], reverse=True)

        # 生成颜色列表，涨用浅红色，跌用浅绿色
        colors = ['lightcoral' if value > 0 else 'lightgreen' for value in draw_stocks_dc.values()]

        # 绘制图
        # 云图部分
        fig = plt.figure(figsize=(12, 16))
        ax1 = fig.add_subplot(211)
        squarify.plot(sizes=draw_sizes,
                      label=label_ls,
                      color=colors, alpha=0.7, pad=True, text_kwargs={'fontsize': 10})

        plt.axis('off')
        plt.title("DDQ选股云图", fontsize=16)
        # 表格部分
        ax2 = fig.add_subplot(212)
        ax2.axis("off")
        table = ax2.table(cellText=table_data[1:], colLabels=table_data[0], loc='center', cellLoc='center')
        table.auto_set_font_size(False)
        table.set_fontsize(10)
        table.scale(1, 1.5)  # 可以调整表格的列宽和行高
        # 修改样式
        for (i, j), cell in table.get_celld().items():
            if i == 0:  # 表头行
                cell.set_facecolor('#D3D3D3')  # 设置表头背景色
                cell.set_text_props(weight='bold', color='black')  # 设置文字为加粗和白色
        # 调整子图间的空间
        plt.tight_layout()

        self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
        plt.close()
        # 返回图像的 Markdown 语法
        return f"!['选股云图'](./{base_name})"

    def plan3_fig(self, market_factor_data_df, symbol, result_dir_path, md_file):

        # 因子相关性表，表格（数值）呈现，包含close
        correlation_matrix = market_factor_data_df.corr()

        # 计算剔除极端值的因子数据
        extreme_ratio = 0.02
        trimmed_df: pd.DataFrame = copy.deepcopy(market_factor_data_df)
        factor_ls = []
        for column in trimmed_df.columns:
            if column not in conf.KChart.not_factor_ls.value:
                q = trimmed_df[column].quantile([extreme_ratio, 1 - extreme_ratio])
                trimmed_df.loc[(trimmed_df[column] < q[extreme_ratio]), column] = q[extreme_ratio]
                trimmed_df.loc[(trimmed_df[column] > q[1 - extreme_ratio]), column] = q[1 - extreme_ratio]
                factor_ls.append(column)

        # 剔除x%极值后的因子特征

        mk_code = ''
        factor_num = len(factor_ls)
        # 因子分布图像
        plt.figure(figsize=(15, 10))
        for i, column in enumerate(factor_ls):
            plt.subplot(factor_num, 1, i + 1)
            draw_sr = market_factor_data_df[column].reset_index(drop=True)
            draw_sr.plot()
            plt.title(f'{symbol}-{column}-自然分布图')
        plt.tight_layout()
        # 保存图像为图片文件
        base_name = f"{symbol}-Density.jpg"
        fig_path = os.path.join(result_dir_path, base_name)

        self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
        mk_code += f"![{symbol} {base_name}](./{base_name})\n---"

        # 因子频率分布直方图（排除x%异常值）
        plt.figure(figsize=(15, 10))
        for i, column in enumerate(factor_ls):
            plt.subplot(factor_num, 1, i + 1)
            market_factor_data_df[column].plot(kind='hist', bins=30, alpha=0.5, label='全数据')
            draw_sr = trimmed_df[column].reset_index(drop=True)
            draw_sr.plot(kind='hist', bins=30, alpha=0.5, label='首尾截断数据')
            plt.title(f'{symbol}-{column}-频数分布图')
            plt.legend()
        plt.tight_layout()
        # 保存图像为图片文件
        base_name = f"{symbol}-Frequency.jpg"
        fig_path = os.path.join(result_dir_path, base_name)

        self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
        mk_code += f"![{symbol} {base_name}](./{base_name})\n---"

        # 异常值分布直方图
        diff_mask = trimmed_df != market_factor_data_df
        outliers_df = market_factor_data_df[diff_mask]
        plt.figure(figsize=(15, 5))
        for i, column in enumerate(factor_ls):
            plt.subplot(factor_num, 1, i + 1)
            draw_sr = outliers_df[column].reset_index(drop=True)
            draw_sr.plot(kind='hist', bins=30)
            plt.title(f'{symbol}-{column}-异常值分布图')
        plt.tight_layout()
        # 保存图像为图片文件
        base_name = f"{symbol}-Outliers.jpg"
        fig_path = os.path.join(result_dir_path, base_name)

        self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
        mk_code += f"![{symbol} {base_name}](./{base_name})\n---"
        # 绘制因子相关性表
        plt.figure(figsize=(15, 5))
        sns.heatmap(correlation_matrix,
                    xticklabels=correlation_matrix.columns.values,
                    yticklabels=correlation_matrix.columns.values,
                    annot=True,  # 显示数值
                    cmap='coolwarm')  # 设置颜色映射
        base_name = f"{symbol}-corr-heatmap.jpg"
        fig_path = os.path.join(result_dir_path, base_name)

        self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
        mk_code += f"![{symbol} {base_name}](./{base_name})\n---"
        md_file.write(mk_code)

        # 绘制全量数据特征描述表格
        # 每个因子的特征统计
        describe_df = market_factor_data_df[factor_ls].describe().round(4)
        tab_str_ls = []
        row = 0
        tab_head = ["--"] + describe_df.columns.tolist()
        tab_str_ls += tab_head
        row += 1
        for index_o in describe_df.index:
            tab_row = [index_o] + describe_df.loc[index_o].astype(str).tolist()
            tab_str_ls += tab_row
            row += 1
        md_file.new_line()
        md_file.new_table(int(len(tab_str_ls) / row), row, tab_str_ls)
        md_file.write("\n---")
        # 绘制剔除极值后特征描述表格
        describe_df = trimmed_df[factor_ls].describe().round(4)
        tab_str_ls = []
        row = 0
        tab_head = ["--"] + describe_df.columns.tolist()
        tab_str_ls += tab_head
        row += 1
        for index_o in describe_df.index:
            tab_row = [index_o] + describe_df.loc[index_o].astype(str).tolist()
            tab_str_ls += tab_row
            row += 1
        md_file.new_line()
        md_file.new_table(int(len(tab_str_ls) / row), row, tab_str_ls)
        md_file.write("\n---")

    def plan_distribution_fig(self, dataframe, plot_type, factor_ls, result_dir_path, symbol):
        dataframe = dataframe.dropna()

        # plan 4 对应的绘图函数，将返回Markdown语法
        if plot_type == "close":
            column = "close"

            plt.figure(figsize=(15, 6))
            base_name = f"{symbol}-{column}-{'收盘价分布图'}.jpg"
            fig_path = os.path.join(result_dir_path, base_name)
            plt.title("收盘价分布图")
            plt.xlabel("时间")
            plt.ylabel("收盘价")

            scale_y = dataframe[column]
            plt.plot(range(len(scale_y.index)), scale_y.values, color="gray",
                     label=f'{column}')
            self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
            plt.close()
            # 返回图像的 Markdown 语法
            return f"![{symbol} {plot_type}](./{base_name})"

        elif plot_type == "factor continuous":
            mk_code = ''
            for i, column in enumerate(factor_ls):
                plt.figure(figsize=(15, 6))
                base_name = f"{symbol}-{column}-{'分布图（连续）'}.jpg"
                fig_path = os.path.join(result_dir_path, base_name)
                plt.title(f"因子{column}分布图（连续）")
                plt.xlabel("时间")
                plt.ylabel("因子值")
                scale_y = dataframe[column]
                plt.plot(range(len(scale_y.index)), scale_y.values, color=plt.cm.tab20(i / len(factor_ls)),
                         label=f'{column}')
                self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
                plt.close()
                # 返回图像的 Markdown 语法
                mk_code += f"![{symbol} {plot_type}](./{base_name})\n---"
            return mk_code
        elif plot_type == "factor dispersed":
            mk_code = ''
            for i, column in enumerate(factor_ls):
                plt.figure(figsize=(15, 6))
                base_name = f"{symbol}-{column}-{'分布图（离散）'}.jpg"
                fig_path = os.path.join(result_dir_path, base_name)
                plt.title(f"因子{column}分布图（离散）")
                plt.xlabel("时间")
                plt.ylabel("因子值")
                scale_y = dataframe[column]
                plt.scatter(range(len(scale_y.index)), scale_y.values, color=plt.cm.tab20(i / len(factor_ls)),
                            label=f'{column}', s=2)
                self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
                plt.close()
                # 返回图像的 Markdown 语法
                mk_code += f"![{symbol} {plot_type}](./{base_name})\n---"
            return mk_code
        elif plot_type == "factor dispersed frequency":
            mk_code = ''
            for i, column in enumerate(factor_ls):
                plt.figure(figsize=(15, 6))
                base_name = f"{symbol}-{column}-{'频数分布直方图（离散）'}.jpg"
                fig_path = os.path.join(result_dir_path, base_name)
                plt.title(f"因子{column}频数分布直方图（离散）")
                plt.xlabel("频数")
                plt.ylabel("因子值")
                frequency = dataframe[column].value_counts()
                bars = plt.bar(frequency.index, frequency.values, color='blue')  # 绘制条形图
                for bar in bars:
                    y_val = bar.get_height()
                    plt.text(bar.get_x() + bar.get_width() / 2, y_val, str(y_val), va='bottom', ha='center')

                self.parent.master.file_manager.save_mpl_fig(plt, fig_path)
                plt.close()
                # 返回图像的 Markdown 语法
                mk_code += f"![{symbol} {plot_type}](./{base_name})\n---"
            return mk_code
        else:
            pass
